import torch
import torch.nn as nn
import torch.nn.functional as F


class LocalContrastiveLossHard(nn.Module):
    def __init__(self, temperature: float = 0.1):
        super().__init__()
        self.temperature = temperature 
        self.cross_entropy = nn.CrossEntropyLoss(reduction='mean')
        
    def forward(self, A, B):
        batch_size, time_steps, hidden_dim = A.shape

        # Normalize embeddings
        A_norm = F.normalize(A, p=2, dim=-1).reshape(batch_size * time_steps, hidden_dim)
        B_norm = F.normalize(B, p=2, dim=-1).reshape(batch_size * time_steps, hidden_dim)

        logits = torch.matmul(A_norm, B_norm.T) / self.temperature

        labels = torch.arange(batch_size * time_steps, device=A.device)

        loss_AB = self.cross_entropy(logits, labels)
        loss_BA = self.cross_entropy(logits.T, labels)

        return (loss_AB + loss_BA) * 0.5
